#include "DBDetect.h"
#include "session.h"
#include "MySQL_proto.h"
#include "tcp_packet.h"
#include "log/log.h"
#include <functional>
#include <fstream>


namespace dbdetect
{
    std::string ProtocolNameString[]{
        "UKNOWN",
        "MYSQL",
        "MARIADB",
        "GBASE8A",
        "TIDB"
    };

    MySQLProtocol::MySQLProtocol(std::string& targetFileName, bool checkCompress, int32_t maxth)
    {
        proto_type_ = PROTOCOL_MYSQL;
        targetFileName_ = targetFileName;
        checkCompress_ = checkCompress;
        maxThreshold_ = maxth;
    }

    int32_t MySQLProtocol::dissector(const uint8_t *buf, size_t size)
    {
        if (size <= 4)
        {
            return RET_UNKNOWN;
        }

        /*if (proto_type_ == PROTOCOL_GBASE8A || proto_type_ == PROTOCOL_TIDB)
        {
            return RET_OVER;
        }*/

        size_t pos = 0;
        if (left_length_ > 0)
        {
            pos += left_length_;
            left_length_ = 0;
        }

        size_t packet_len = ((uint8_t)(buf[pos + 2]) << 16) + ((uint8_t)(buf[pos + 1]) << 8) + (uint8_t)(buf[pos]) + 4;
        if (packet_len == size)
        {
            uint8_t packet_number = (uint8_t)buf[3];
            uint8_t server_version = (uint8_t)buf[5];
            if (packet_number == 0 && server_version > 0x30 && server_version < 0x39 && buf[6] == 0x2e)
            {
                size_t pos2 = 5;
                do
                {
                    while (pos2 < size && buf[pos2] != 0x54)
                    { // T
                        pos2++;
                    }
                    pos2++;

                    if (pos2 >= size)
                    {
                        break;
                    }

                    if (buf[pos2] != 0x69)
                    { // i
                        break;
                    }
                    pos2++;

                    if (buf[pos2] != 0x44)
                    { // D
                        break;
                    }
                    pos2++;

                    if (buf[pos2] != 0x42)
                    { // B
                        break;
                    }
                    LOG_DEBUG("TiDB : Server Greeting detected!");
                    proto_type_ = PROTOCOL_TIDB;
                    return RET_S2C;
                } while (0);

                phase_ = PHASE_FULL_LINK;
                set_proto_type();
                set_proto_version((const char*)buf+5);          
                return RET_S2C;
            }

            // Login Request Packet
            if (phase_ == PHASE_FULL_LINK && packet_number == 1 && !gotLoginReq_)
            {
                LOG_DEBUG("MySQL : Login Request Packet detected!");
                set_proto_type();
                gotLoginReq_ = true;
                return RET_C2S;
            }

            uint8_t flag = (uint8_t)buf[4];

            // Response Packet
            if (flag == 0 || flag == 0xFE || flag == 0xFF)
            {
                LOG_DEBUG("MySQL : Response Packet detected!");
                set_proto_type();
                return RET_S2C;
            }

            // Command Packet
            if (flag == 1 && packet_len == 1)
            {
                // COM_QUIT
                LOG_DEBUG("MySQL : Command Packet detected!");
                set_proto_type();
                return RET_C2S;
            }

            // Query Packet
            if (flag == 3)
            {
                // 检查是否为gbase
                if (buf[8] == 0x67 && buf[9] == 0x62 && buf[10] == 0x61 && buf[11] == 0x73 && buf[12] == 0x65)
                {
                    LOG_DEBUG("GBASE8A : Query Packet detected!");
                    proto_type_ = PROTOCOL_GBASE8A;
                    return RET_C2S;
                }
                LOG_DEBUG("MySQL : Query Packet detected!");
            }

            if (flag > 1 && flag < 0x1f)
            {
                set_proto_type();
                return RET_C2S;
            }
        }
        else if (packet_len < size)
        {
            //uint32_t first_packet_number = buf[pos + 3];
            uint32_t last_packet_number = buf[pos + 3];

            // Check TABULAR Resp
            pos += packet_len;

            size_t total_len = pos;
            size_t nLen = size - total_len;

            uint32_t packet_number = packet_number_ + 1;
            while (nLen > 0 && pos + 4 < size)
            {
                size_t pdu_size = ((uint8_t)(buf[pos + 2]) << 16) + ((uint8_t)(buf[pos + 1]) << 8) + (uint8_t)(buf[pos]) + 4;
                last_packet_number = buf[pos + 3];
                if (pdu_size > size)
                {
                    break;
                }
                nLen -= pdu_size;
                pos += pdu_size;
                total_len += pdu_size;

                packet_number = (packet_number == 255) ? 0 : packet_number + 1;
            }

            if (total_len == size)
            {
                left_length_ = 0;
                packet_number_ = last_packet_number;
                LOG_DEBUG("MySQL : TABULAR Resp detected!");
                return RET_S2C;
            }

            // Packet data segment,Check Packet Number
            if (total_len > size)
            {
                left_length_ = total_len - size;
                if (packet_number == last_packet_number)
                {
                    LOG_DEBUG("MySQL : Packet data segment detected!");
                    packet_number_ = last_packet_number;
                    return RET_S2C;
                }
            }

            left_length_ = 0;
            packet_number_ = 0;

            if (checkCompress_ && packet_len+3 == size)
            {
                if (checkIfTheDataIsCompressed(buf))
                {
                    LOG_DEBUG("MySQL : Compressed Packet detected!");
                    set_proto_type();
                    if (last_packet_number == 0)
                    {
                        return RET_C2S;
                    }
                    else if (last_packet_number == 1)
                    {
                        return RET_S2C;
                    }
                    return RET_UNSURE;
                }
            }
        }

        return RET_UNKNOWN;
    }

    void MySQLProtocol::set_proto_type()
    {
        proto_type_ = PROTOCOL_MYSQL;
    }

    void MySQLProtocol::set_proto_version(const char* ver)
    {        
        proto_version_ = ver;
        auto pos = proto_version_.rfind("MariaDB");
        if (proto_type_ == PROTOCOL_MYSQL && pos !=std::string::npos)
        {
            proto_type_ = PROTOCOL_MARIADB;
            proto_version_ = proto_version_.substr(0, pos-1);
            LOG_DEBUG("MariaDB : Server Greeting detected!");
        }
        else{
            LOG_DEBUG("MySQL : Server Greeting detected!");
        }
    }

    size_t hashIpQuadruple(uint32_t src_ip, uint32_t dst_ip, uint16_t src_port, uint16_t dst_port)
    {
        std::hash<uint32_t> hash_uint32;
        std::hash<uint16_t> hash_uint16;

        size_t src_ip_hash = hash_uint32(src_ip);
        size_t dst_ip_hash = hash_uint32(dst_ip);
        size_t src_port_hash = hash_uint16(src_port);
        size_t dst_port_hash = hash_uint16(dst_port);

        // 使用XOR组合各个哈希值
        return src_ip_hash ^ dst_ip_hash ^ src_port_hash ^ dst_port_hash;
    }

    void MySQLProtocol::setHashDBTypeMap(size_t socketHash){
        if (proto_type_ != PROTOCOL_UNKNOWN)
            parsed_DB_map_[socketHash] = proto_type_;
    }

    bool MySQLProtocol::isSocketInMap(size_t socketHash){
        return parsed_DB_map_.find(socketHash) != parsed_DB_map_.end();
    }    

    bool MySQLProtocol::writeTargetFile(secsmart_tcpip::IPAddress& sip, secsmart_tcpip::IPAddress& dip, uint16_t sport, uint16_t dport, bool s2c)
    {
        if (targetFileName_.empty())
        {
            return false;
        }
        std::string strOut = "{ \"DB_Type\":\"" +ProtocolNameString[proto_type_] + "\", ";
        strOut += "\"DB_Version\":\"" +proto_version_ + "\", ";
        secsmart_tcpip::IPAddress ip;
        uint16_t port=0;
        if (s2c)
        {
            ip = sip;
            port = sport; 
        }
        else
        {
            ip = dip;
            port = dport;
        }
        /*if (port > 10000)
        {
            if(s2c && sport > dport)
            {
                ip = dip;
                port = dport;
            }
            else if (!s2c && dport > sport)
            {
                ip = sip;
                port = sport;
            }
        }*/
        strOut += "\"ip\":\""+ ip.toString() + "\", ";
        strOut += "\"port\":\""+ std::to_string(port) + "\" }";

        std::ofstream ofs(targetFileName_, std::ios_base::app);
        ofs << strOut << std::endl;
        LOG_DEBUG("{}", strOut);

        return true;
    }

    bool MySQLProtocol::checkIfTheDataIsCompressed(const uint8_t *data)
    {
        unsigned long tmp_pktLen = 0;
        unsigned long tmp_compBufLen = (data[2] << 16) | (data[1] << 8) | data[0];
        if (0 == data[4] && 0 == data[5] && 0 == data[6])
        {
            // 压缩协议开启，但数据未进行压缩
            tmp_pktLen = (data[9] << 16) | (data[8] << 8) | (data[7]);
        }
        return (tmp_compBufLen == (tmp_pktLen + 4));
    }

    std::shared_ptr<MySQLProtocol> initParser(std::string &targetFileName, int32_t maxth,bool checkCompress)
    {
        auto ptr = std::make_shared<MySQLProtocol>(targetFileName, checkCompress, maxth);
        return ptr;
    }

    void parseData(secsmart_session::session_sptr pSession, std::shared_ptr<MySQLProtocol> parserPTR, const secsmart_tcpip::tcp_stream_data_t *data)
    {
        if (!parserPTR)
        {
            return;
        }

        size_t sh = hashIpQuadruple(pSession->sip.getHash(), pSession->dip.getHash(), pSession->sport, pSession->dport);
        if (parserPTR->isSocketInMap(sh))
        {
            pSession->dbType = parserPTR->get_proto_type();
            //pSession->detectSatus = secsmart_session::DATABASE_DETECTION_OK;
            return;
        }
        LOG_DEBUG("dissector: {}:{} {} {}:{} dataLen={}", pSession->sip.toString(), pSession->sport, data->is_forward?"-->":"<--", pSession->dip.toString(), pSession->dport, data->data_len);
        auto ret = parserPTR->dissector(data->data, data->data_len);
        if (ret == RET_C2S || ret ==RET_S2C)
        {
            pSession->matchCount++;
            LOG_DEBUG("DB Detected, Ret={}, Type={}, Match count={}, MAX={} ", ret, parserPTR->get_proto_type(), pSession->matchCount, parserPTR->getMaxThreshold());
        }
        else if (ret == RET_UNKNOWN)
        {
            pSession->unknowPktCount++;
            if (pSession->unknowPktCount > 10) { // 未知报文数大于10个时，不解析后续报文
                pSession->dbType = PROTOCOL_UNKNOWN;
                pSession->detectSatus = secsmart_session::DATABASE_DETECTION_UNKNOWN;
                parserPTR->setHashDBTypeMap(sh);
            }
            return;
        }

        if (pSession->matchCount >= parserPTR->getMaxThreshold())
        {
            pSession->dbType = parserPTR->get_proto_type();
            pSession->dbVersion = parserPTR->get_proto_version();
            /*if (ret == RET_C2S)
            {
                pSession->direction = secsmart_session::SESSION_DIRCTION_REQUEST;
            }
            else if(ret == RET_S2C)
            {
                pSession->direction = secsmart_session::SESSION_DIRCTION_RESPONSE;
            }
            else if(ret == RET_UNSURE)
            {
                pSession->direction = secsmart_session::SESSION_DIRCTION_UNKNOW;
            }*/
            pSession->detectSatus = secsmart_session::DATABASE_DETECTION_OK;
            parserPTR->setHashDBTypeMap(sh);
            parserPTR->writeTargetFile(pSession->sip, pSession->dip, pSession->sport, pSession->dport, ret==RET_C2S ^ data->is_forward);
            //LOG_INFO("DB Detected, type {}  RET = {} is_forward = {} match={} MAX={}", pSession->dbType , ret, data->is_forward,pSession->matchCount, parserPTR->getMaxThreshold());
            return;
        }
    }
}; // namespace dbdtect